#include "exception.h"

#include <stdlib.h>
#include <assert.h>
#include "memory.h"

#define NUM_FRAMES_PER_ALLOCATION   64


typedef struct ExceptionContext ExceptionContext;


/**
 * The different states of an ExceptionFrame (e.g. an exception has been thrown, it has been caught...)
 */
typedef enum ExceptionState
{
  ExceptionState_None    = 0x00,
  ExceptionState_Thrown  = 0x01,
  ExceptionState_Caught  = 0x02
} ExceptionState;


/**
 * A context is a stack of exception-frames. Each time the program enters a "try/catch" block, an
 * exception-frame is pushed. Each time the program leaves a "try/catch" block, the top exception-frame is
 * popped. In multi-threaded programs, there should be one context per thread.
 */
struct ExceptionContext
{
  ExceptionFrame* frames;
  ExceptionFrame* topFrame;
  int numFrames;
  int numAllocatedFrames;
};


// TODO: use per-thread storage to make exception-handling works with several threads
static ExceptionContext s_context = { NULL, NULL, 0, 0 };


/**
 * Create a new exception-frame and push it on top of the context. When an exception-handling block begins,
 * this function has to be called and setjmp has to be used on the jmp-buffer of the returned exception-frame,
 * in order to save the context to restore when an exception is thrown.
 * exceptionBegin() is a convenience macro that does exactly that.
 */
ExceptionFrame* exceptionNewFrame(void)
{
  ExceptionContext* context = &s_context;

  // Allocate new exception-frames if needed
  if (context->numFrames == context->numAllocatedFrames) {
    context->numAllocatedFrames += NUM_FRAMES_PER_ALLOCATION;
    context->frames = memoryRealloc(context->frames, context->numAllocatedFrames * sizeof(ExceptionFrame));
  }

  // Set-up and push the new exception-frame
  ExceptionFrame* newFrame = &context->frames[context->numFrames];
  newFrame->prevFrame = context->topFrame;
  newFrame->exception = NULL;
  newFrame->block = ExceptionBlock_Try;
  newFrame->state = ExceptionState_None;
  context->topFrame = newFrame;
  context->numFrames++;

  return newFrame;
}


/**
 * Throw the given exception. When called from a try block, it jumps to the corresponding catch block.
 * Otherwise, the finally block is executed if it hasn't been yet, and the exception is propagated to the
 * previous exception-frame
 */
void exceptionThrow(void* exception)
{
  ExceptionContext* context = &s_context;
  ExceptionFrame* curFrame = context->topFrame;
  assert(curFrame);

  // Store the exception in the current frame
  curFrame->exception = exception;
  curFrame->state |= ExceptionState_Thrown;

  // If the exception has been thrown from the try block, jump to the catch block
  if (curFrame->block == ExceptionBlock_Try)
  {
    curFrame->block = ExceptionBlock_Catch;
    longjmp(curFrame->jmpBuf, curFrame->block);
  }

  // If the exception has been thrown from the catch block, jump to the finally block
  else if (curFrame->block == ExceptionBlock_Catch) {
    curFrame->state &= ~ExceptionState_Caught;
    curFrame->block = ExceptionBlock_Finally;
    longjmp(curFrame->jmpBuf, curFrame->block);
  }

  // Otherwise, if the exception has been thrown from the finally block. In this case, let's leave
  // this exception-handling block and jump to the previous frame
  else
  {
    curFrame->state &= ~ExceptionState_Caught;
    exceptionEnd();
  }
}


/**
 * This function has to be called from a catch block. It returns the exception that has been thrown in the
 * corresponding try block (same as exceptionCatch() except it does not mark the exception has caught)
 */
void* exceptionPeek(void)
{
  ExceptionContext* context = &s_context;
  ExceptionFrame* curFrame = context->topFrame;
  assert(curFrame && (curFrame->block == ExceptionBlock_Catch));
  return curFrame->exception;
}


/**
 * This function has to be called from a catch block. It marks as caught the exception thrown in the
 * corresponding try block. This way, the exception won't be propagated to the previous exception-frame
 * when the exception-handling block ends (i.e. when exceptionEnd() is called). The exception is returned.
 */
void* exceptionCatch(void)
{
  ExceptionContext* context = &s_context;
  ExceptionFrame* curFrame = context->topFrame;
  assert(curFrame && (curFrame->block == ExceptionBlock_Catch));
  curFrame->state |= ExceptionState_Caught;
  return curFrame->exception;
}


/**
 * This function has to be called to leave an exception-handling block. It makes sure the finally-block is
 * always called before leaving. If an exception has been thrown but was not caught by the current frame,
 * the exception is thrown to the previous exception-frame.
 */
void exceptionEnd(void)
{
  ExceptionContext* context = &s_context;
  ExceptionFrame* curFrame = context->topFrame;
  assert(curFrame);

  // If the finally-block has not been executed yet, execute it
  if (curFrame->block != ExceptionBlock_Finally)
  {
    curFrame->block = ExceptionBlock_Finally;
    longjmp(curFrame->jmpBuf, curFrame->block);
  }

  // Otherwise, pop the top exception-frame out of the exception-context
  else {
    context->topFrame = curFrame->prevFrame;
    context->numFrames--;

    // If an exception has been thrown but has not been caught, propagate the exception to the previous frame
    if ((curFrame->state & ExceptionState_Thrown) && !(curFrame->state & ExceptionState_Caught))
    {
      exceptionThrow(curFrame->exception);
    }
  }
}
